import numpy as np
import math

# ====== 常量与场景 ======
g = 9.81  # m/s^2
R = 10.0  # 有效遮蔽半径
v_sink = 3  # 起爆后云团中心下沉速度 (m/s)
T = np.array([0, 200, 0])  # 真目标中心
Of = np.array([0, 0, 0])  # 假目标（原点）

# 导弹 M1
M0 = np.array([20000, 0, 2000])
uM = (Of - M0) / np.linalg.norm(Of - M0)
vM = 300 * uM  # 300 m/s 指向假目标
L = np.linalg.norm(Of - M0)
t_hit = L / 300  # 导弹到达假目标时刻≈ 67 s

# FY1（等高度飞行，z=1800 m）
F0 = np.array([17800, 0, 1800])
v_min = 70
v_max = 140

# ====== 搜索设置（可按需调节速度/精度）======
# 航向角与速度的粗网格
angles = np.deg2rad(np.arange(-180, 180, 15))  # 相对 +x 轴，逆时针为正
speeds = np.arange(70, 141, 10)  # m/s

# 三枚弹的起爆时刻扫描（全局时刻）；引信延时候选
dt_eval = 0.02  # 时间评估步长（全局）
t_grid = np.arange(2, t_hit - 2, 2)  # 候选起爆时刻（2 s 步长）
tau_set = np.array([3.0, 3.6, 4.2])  # 候选引信延时（s）
min_release_gap = 1.0  # 相邻两次"投放"时间间隔≥ 1 s
tau_min = 0.5  # 最小引信延时（安全下限）
tau_max = 6.0  # 最大引信延时（便于搜索）

# ====== 预计算导弹轨迹与"导弹→真目标"视线量 ======
ts = np.arange(0, t_hit + dt_eval, dt_eval)  # 全局时间轴
N = len(ts)
Mx = M0[0] + vM[0] * ts
My = M0[1] + vM[1] * ts
Mz = M0[2] + vM[2] * ts

Ux = T[0] - Mx
Uy = T[1] - My
Uz = T[2] - Mz  # 视线向量
UU = Ux ** 2 + Uy ** 2 + Uz ** 2

# ====== 贪心三弹搜索（对每个θ、v 逐弹选最优）======
best_total = -np.inf
best_plan = {}

for ia, theta in enumerate(angles):
    uxF = np.cos(theta)
    uyF = np.sin(theta)

    for iv, vF in enumerate(speeds):
        # 逐弹贪心：初始化（当前联合最小距离为 +Inf）
        dmin_all = np.full(N, np.inf)
        shots = []
        last_t_rel = -np.inf  # 上一次"投放"时间（约束间隔）

        for k in range(3):  # 三枚弹
            best_k_gain = -np.inf
            best_k = {}

            for t_det in t_grid:
                for tau in tau_set:
                    if tau < tau_min or tau > tau_max:
                        continue

                    t_rel = t_det - tau
                    if t_rel < 0:
                        continue

                    if t_rel < last_t_rel + min_release_gap:
                        continue

                    # 起爆点（水平由 FY1 位移给出，竖直由自由落体给出）
                    x_det = F0[0] + vF * uxF * t_det
                    y_det = F0[1] + vF * uyF * t_det
                    z_det = 1800 - 0.5 * g * tau ** 2

                    # 云团中心轨迹（仅在 [t_det, t_det+20] 内有效）
                    Cz = z_det - v_sink * (ts - t_det)  # 数组
                    valid = (ts >= t_det) & (ts <= t_det + 20)

                    # 计算该弹对"导弹→真目标"线段的最小距离 d_k(t)
                    Wx = x_det - Mx
                    Wy = y_det - My
                    Wz = Cz - Mz

                    WU = Wx * Ux + Wy * Uy + Wz * Uz
                    lam = WU / UU
                    lam = np.clip(lam, 0, 1)

                    Px = Mx + lam * Ux
                    Py = My + lam * Uy
                    Pz = Mz + lam * Uz

                    dx = x_det - Px
                    dy = y_det - Py
                    dz = Cz - Pz

                    d_k = np.sqrt(dx ** 2 + dy ** 2 + dz ** 2)
                    d_k[~valid] = np.inf  # 窗口外无效

                    # 与现有联合取最小
                    dmin_new = np.minimum(dmin_all, d_k)
                    mask_new = (dmin_new <= R)
                    gain = np.sum(mask_new) * dt_eval  # 联合遮蔽总时长（秒）

                    if gain > best_k_gain:
                        best_k_gain = gain
                        best_k = {
                            't_det': t_det,
                            'tau': tau,
                            't_rel': t_rel,
                            'x_det': x_det,
                            'y_det': y_det,
                            'z_det': z_det,
                            'dmin_new': dmin_new.copy()
                        }

            # 若该轮未找到可行解，则提前结束
            if not best_k:
                break

            # 接受该弹的最优选择，更新联合 dmin_all
            dmin_all = best_k['dmin_new']
            last_t_rel = best_k['t_rel']

            # 记录该弹信息
            shot_info = {
                't_det': best_k['t_det'],
                'tau': best_k['tau'],
                't_rel': best_k['t_rel'],
                'D': np.array([F0[0] + vF * uxF * best_k['t_rel'],
                               F0[1] + vF * uyF * best_k['t_rel'],
                               1800]),
                'C0': np.array([best_k['x_det'], best_k['y_det'], best_k['z_det']])
            }
            shots.append(shot_info)

        total_cov = np.sum(dmin_all <= R) * dt_eval
        if total_cov > best_total:
            best_total = total_cov
            best_plan = {
                'theta': theta,
                'vF': vF,
                'shots': shots,
                'dmin': dmin_all.copy()
            }

# ====== 输出最优策略 ======
print("========= 问题 3 最优三弹遮蔽策略（FY1→M1） =========")
print(f"航向角θ = {np.rad2deg(best_plan['theta']):.2f}°（相对 +x 逆时针）")
print(f"飞行速度 v = {best_plan['vF']:.2f} m/s")

for k, s in enumerate(best_plan['shots'], 1):
    print(f"——第{k} 枚——")
    print(f"投放时刻 t_rel = {s['t_rel']:.3f} s （与上一枚间隔满足≥1 s）")
    print(f"引信延时τ = {s['tau']:.3f} s")
    print(f"起爆时刻 t_det = {s['t_det']:.3f} s")
    print(f"投放点 D{k} = ({s['D'][0]:.3f}, {s['D'][1]:.3f}, {s['D'][2]:.3f}) m")
    print(f"起爆点 C{k} = ({s['C0'][0]:.3f}, {s['C0'][1]:.3f}, {s['C0'][2]:.3f}) m")

print(f"三弹联合最大遮蔽总时长 = {best_total:.6f} s")

# 理论上单弹≤20 s，三弹≤60 s（若超过 60 s，说明设置/步长需检查）
if best_total > 60 + 1e-6:
    print("警告：遮蔽总时长超过 60 s，建议检查搜索/评估设置。")